import base64
from typing import Optional
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.fernet import Fernet, InvalidToken
from passlib.context import CryptContext
import rsa
import logging

from server.etc.settings import settings

logger = logging.getLogger(__name__)

# 密码哈希上下文
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")


class AESEncryption:
    """AES加密工具"""

    def __init__(self, password: str, salt: bytes = b"salt"):
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        self.cipher = Fernet(key)

    def encrypt(self, data: str) -> str:
        """加密数据"""
        return self.cipher.encrypt(data.encode()).decode()

    def decrypt(self, encrypted_data: str) -> Optional[str]:
        """解密数据"""
        try:
            return self.cipher.decrypt(encrypted_data.encode()).decode()
        except InvalidToken:
            logger.error("Failed to decrypt data: invalid token")
            return None


class RSAEncryption:
    """RSA加密工具"""

    def __init__(self):
        # 在实际应用中，应该从配置或安全存储中加载密钥
        self.public_key, self.private_key = self._generate_keys()

    def _generate_keys(self, key_size: int = 2048):
        """生成RSA密钥对"""
        return rsa.newkeys(key_size)

    def encrypt(self, data: str) -> str:
        """RSA加密"""
        encrypted = rsa.encrypt(data.encode(), self.public_key)
        return base64.b64encode(encrypted).decode()

    def decrypt(self, encrypted_data: str) -> Optional[str]:
        """RSA解密"""
        try:
            decoded = base64.b64decode(encrypted_data.encode())
            return rsa.decrypt(decoded, self.private_key).decode()
        except Exception as e:
            logger.error(f"Failed to decrypt RSA data: {str(e)}")
            return None


# 全局加密工具实例
aes_encryptor = AESEncryption(settings.JWT_SECRET_KEY)
rsa_encryptor = RSAEncryption()


def encrypt_field(value: str) -> str:
    """加密敏感字段"""
    if not value:
        return value
    return aes_encryptor.encrypt(value)


def decrypt_field(encrypted_value: str) -> Optional[str]:
    """解密敏感字段"""
    if not encrypted_value:
        return encrypted_value
    return aes_encryptor.decrypt(encrypted_value)